//
// Created by corey on 2021/4/13.
//

#ifndef CORNETNET_LITE_TRT_BOTTOM_POOL_PLUGIN_H
#define CORNETNET_LITE_TRT_BOTTOM_POOL_PLUGIN_H

#include "corner_pool_plugin.h"

class BottomPoolPlugin : public CornerPoolPlugin
{
public:
    BottomPoolPlugin() : CornerPoolPlugin()
    {}

    BottomPoolPlugin(const void *buffer, size_t length) : CornerPoolPlugin(buffer, length)
    {}

    IPluginV2Ext *clone() const TRTNOEXCEPT override
    {
        return new BottomPoolPlugin(&m, sizeof(m));
    }

    int enqueue(int batchSize, const void *const *inputs, void **outputs, void *workspace,
                cudaStream_t stream) TRTNOEXCEPT override
    {
        BottomPooling(inputs[0], outputs[0],
                      batchSize,
                      m.input_dim.d[0],
                      m.input_dim.d[1],
                      m.input_dim.d[2],
                      m.data_type,
                      stream);
        return 0;
    }

    const char *getPluginType() const override
    { return "BottomPool"; }
};

class BottomPoolPluginCreator : public nvinfer1::IPluginCreator
{
public:
    static nvinfer1::PluginFieldCollection plugin_field_collection;

    nvinfer1::IPluginV2 *deserializePlugin(const char *name, const void *serialData, size_t serialLength) override
    {
        return new BottomPoolPlugin(serialData, serialLength);
    }

    const char *getPluginName() const override
    { return "BottomPool"; }

    const char *getPluginVersion() const override
    { return "1"; }

    void setPluginNamespace(const char *szNamespace) override
    {}

    const char *getPluginNamespace() const override
    { return ""; }

    const nvinfer1::PluginFieldCollection *getFieldNames() override
    {
        return &plugin_field_collection;
    }

    nvinfer1::IPluginV2 *createPlugin(const char *name, const nvinfer1::PluginFieldCollection *fc) override
    {
        return new BottomPoolPlugin();
    }
};

nvinfer1::PluginFieldCollection BottomPoolPluginCreator::plugin_field_collection = nvinfer1::PluginFieldCollection();

REGISTER_TENSORRT_PLUGIN(BottomPoolPluginCreator);

#endif
